#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
视觉分析模块主入口
整合视频处理和视觉大模型分析功能
"""

import os
import sys
import logging
from typing import Dict, Any, Optional, List, Tuple
from datetime import datetime

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('vision_analysis.log', encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)

# 导入子模块
from src.vision.video_processor import VideoProcessor
from src.vision.zhipu_vision_analyzer import ZhipuVisionAnalyzer
from src.vision.analysis_result_manager import AnalysisResultManager

# 设置默认配置
DEFAULT_RESULT_DIR = "./analysis_results"
MAX_KEYFRAMES = 5


class VisionAnalysisModule:
    """
    视觉分析模块
    整合视频处理、AI分析和结果管理功能
    """
    
    def __init__(self, 
                 api_key: str = None,
                 api_url: str = None,
                 result_dir: str = DEFAULT_RESULT_DIR,
                 max_keyframes: int = MAX_KEYFRAMES,
                 config_path: str = None):
        """
        初始化视觉分析模块
        
        Args:
            api_key: 智谱API密钥，如果为None则从环境变量或配置文件加载
            api_url: 智谱API地址
            result_dir: 结果存储目录
            max_keyframes: 最大关键帧数量
            config_path: 配置文件路径
        """
        # 初始化组件
        self.video_processor = VideoProcessor()
        self.result_manager = AnalysisResultManager(result_dir)
        self.analyzer = ZhipuVisionAnalyzer(
            api_key=api_key,
            api_url=api_url,
            config_path=config_path
        )
        self.max_keyframes = max_keyframes
        
        logger.info("视觉分析模块初始化完成")
    
    def analyze_video(self, 
                     video_path: str, 
                     video_id: str = None,
                     force_reanalyze: bool = False,
                     keyframe_strategy: str = "content_based",
                     encoding_categories: List[str] = None) -> Dict[str, Any]:
        """
        分析视频
        
        Args:
            video_path: 视频文件路径
            video_id: 视频ID，如果为None则根据文件名生成
            force_reanalyze: 是否强制重新分析
            keyframe_strategy: 关键帧提取策略 (content_based/uniform)
            encoding_categories: 要分析的编码类目列表，如果为None则分析所有类目
            
        Returns:
            Dict[str, Any]: 分析结果
        """
        # 验证视频文件
        if not os.path.exists(video_path):
            logger.error(f"视频文件不存在: {video_path}")
            raise FileNotFoundError(f"视频文件不存在: {video_path}")
        
        # 生成或使用视频ID
        if video_id is None:
            video_id = self._generate_video_id(video_path)
        
        # 检查是否已存在分析结果
        if not force_reanalyze:
            existing_result = self.result_manager.load_analysis_result(video_id)
            if existing_result:
                logger.info(f"已存在分析结果，直接返回: {video_id}")
                return existing_result
        
        # 开始分析流程
        logger.info(f"开始分析视频: {video_path} (ID: {video_id})")
        
        try:
            # 步骤1: 提取视频信息
            video_info = self.video_processor.get_video_info(video_path)
            logger.info(f"视频信息: 时长={video_info['duration']:.2f}秒, 分辨率={video_info['width']}x{video_info['height']}")
            
            # 步骤2: 提取关键帧
            keyframes = self.video_processor.extract_keyframes(
                video_path,
                max_frames=self.max_keyframes,
                strategy=keyframe_strategy
            )
            logger.info(f"成功提取 {len(keyframes)} 个关键帧")
            
            # 步骤3: 调用视觉大模型分析
            analysis_result = self.analyzer.analyze_video_frames(
                keyframes=keyframes,
                video_id=video_id,
                video_info=video_info,
                encoding_categories=encoding_categories
            )
            
            # 步骤4: 添加视频信息到结果
            analysis_result['video_path'] = video_path
            analysis_result['video_info'] = video_info
            analysis_result['keyframe_count'] = len(keyframes)
            analysis_result['analysis_time'] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            
            # 步骤5: 计算总体置信度
            analysis_result['overall_confidence'] = self._calculate_overall_confidence(analysis_result)
            
            # 步骤6: 保存结果
            self.result_manager.save_analysis_result(video_id, analysis_result)
            
            logger.info(f"视频分析完成: {video_id}")
            return analysis_result
            
        except Exception as e:
            logger.error(f"视频分析失败: {str(e)}", exc_info=True)
            
            # 保存失败状态
            error_result = {
                "video_id": video_id,
                "video_path": video_path,
                "analysis_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "status": "failed",
                "error_message": str(e)
            }
            self.result_manager.save_analysis_result(video_id, error_result, status="failed")
            
            raise
        
        finally:
            # 清理临时文件
            if 'keyframes' in locals():
                for frame_path, _ in keyframes:
                    if os.path.exists(frame_path):
                        try:
                            os.remove(frame_path)
                        except:
                            pass
    
    def batch_analyze_videos(self, 
                           video_paths: List[str],
                           parallel: bool = False,
                           max_workers: int = 4,
                           **kwargs) -> Dict[str, Dict[str, Any]]:
        """
        批量分析视频
        
        Args:
            video_paths: 视频文件路径列表
            parallel: 是否并行分析
            max_workers: 最大工作线程数
            **kwargs: 传递给analyze_video的其他参数
            
        Returns:
            Dict[str, Dict[str, Any]]: {video_path: analysis_result}
        """
        results = {}
        
        if parallel and len(video_paths) > 1:
            # 并行分析
            from concurrent.futures import ThreadPoolExecutor, as_completed
            
            with ThreadPoolExecutor(max_workers=min(max_workers, len(video_paths))) as executor:
                future_to_video = {executor.submit(self.analyze_video, path, **kwargs): path for path in video_paths}
                
                for future in as_completed(future_to_video):
                    video_path = future_to_video[future]
                    try:
                        result = future.result()
                        results[video_path] = result
                        logger.info(f"批量分析完成: {video_path}")
                    except Exception as e:
                        logger.error(f"批量分析失败: {video_path}, 错误: {str(e)}")
                        results[video_path] = {
                            "status": "failed",
                            "error_message": str(e)
                        }
        else:
            # 串行分析
            for video_path in video_paths:
                try:
                    result = self.analyze_video(video_path, **kwargs)
                    results[video_path] = result
                except Exception as e:
                    logger.error(f"批量分析失败: {video_path}, 错误: {str(e)}")
                    results[video_path] = {
                        "status": "failed",
                        "error_message": str(e)
                    }
        
        return results
    
    def validate_analysis_result(self, result: Dict[str, Any]) -> Dict[str, Any]:
        """
        验证分析结果
        
        Args:
            result: 分析结果
            
        Returns:
            Dict[str, Any]: 验证报告
        """
        validation = {
            "is_valid": True,
            "missing_fields": [],
            "low_confidence_fields": [],
            "errors": []
        }
        
        # 检查必需字段
        required_fields = ["video_id", "analysis_time", "status"]
        for field in required_fields:
            if field not in result:
                validation["missing_fields"].append(field)
                validation["is_valid"] = False
        
        # 检查编码类目
        if "encoding_categories" in result:
            for category_name, category_data in result["encoding_categories"].items():
                # 检查每个类目的必需字段
                if "option" not in category_data or "confidence" not in category_data:
                    validation["missing_fields"].append(f"{category_name}: option/confidence")
                    validation["is_valid"] = False
                
                # 检查置信度
                if "confidence" in category_data and category_data["confidence"] < 50:
                    validation["low_confidence_fields"].append(f"{category_name}: {category_data['confidence']}%")
        
        return validation
    
    def generate_analysis_report(self, 
                               video_id: str,
                               output_format: str = "json",
                               output_path: str = None) -> str:
        """
        生成分析报告
        
        Args:
            video_id: 视频ID
            output_format: 输出格式 (json/text)
            output_path: 输出路径
            
        Returns:
            str: 报告内容
        """
        # 加载分析结果
        result = self.result_manager.load_analysis_result(video_id)
        if not result:
            raise ValueError(f"未找到视频分析结果: {video_id}")
        
        # 生成报告
        if output_format == "text":
            report = self._generate_text_report(result)
        else:
            report = json.dumps(result, ensure_ascii=False, indent=2)
        
        # 保存报告
        if output_path:
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(report)
            logger.info(f"分析报告已保存: {output_path}")
        
        return report
    
    def export_all_results(self, 
                         output_dir: str,
                         format: str = "json") -> List[str]:
        """
        导出所有分析结果
        
        Args:
            output_dir: 输出目录
            format: 输出格式
            
        Returns:
            List[str]: 导出文件路径列表
        """
        # 确保输出目录存在
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        exported_files = []
        
        # 获取所有结果
        all_results = self.result_manager.list_all_results()
        
        for result_summary in all_results:
            video_id = result_summary["video_id"]
            result = self.result_manager.load_analysis_result(video_id)
            
            if result:
                # 生成输出路径
                output_path = os.path.join(output_dir, f"analysis_{video_id}.{format}")
                
                # 保存结果
                try:
                    if format == "text":
                        with open(output_path, 'w', encoding='utf-8') as f:
                            f.write(self._generate_text_report(result))
                    else:
                        with open(output_path, 'w', encoding='utf-8') as f:
                            json.dump(result, f, ensure_ascii=False, indent=2)
                    
                    exported_files.append(output_path)
                except Exception as e:
                    logger.error(f"导出失败: {video_id}, 错误: {str(e)}")
        
        logger.info(f"成功导出 {len(exported_files)} 个分析结果")
        return exported_files
    
    def _generate_video_id(self, video_path: str) -> str:
        """
        根据视频路径生成视频ID
        
        Args:
            video_path: 视频文件路径
            
        Returns:
            str: 视频ID
        """
        import hashlib
        
        # 使用文件名和文件大小生成唯一ID
        file_name = os.path.basename(video_path)
        file_size = os.path.getsize(video_path)
        
        # 创建ID
        unique_str = f"{file_name}_{file_size}"
        video_id = hashlib.md5(unique_str.encode()).hexdigest()[:10]
        
        # 添加时间戳前缀以确保唯一性
        timestamp = datetime.now().strftime("%Y%m%d")
        
        return f"{timestamp}_{video_id}"
    
    def _calculate_overall_confidence(self, result: Dict[str, Any]) -> float:
        """
        计算总体置信度
        
        Args:
            result: 分析结果
            
        Returns:
            float: 总体置信度
        """
        confidences = []
        
        # 从编码类目中提取置信度
        if "encoding_categories" in result:
            for category_data in result["encoding_categories"].values():
                if "confidence" in category_data:
                    confidences.append(category_data["confidence"])
        
        # 如果没有置信度，返回默认值
        if not confidences:
            return 0.0
        
        # 计算平均值
        return sum(confidences) / len(confidences)
    
    def _generate_text_report(self, result: Dict[str, Any]) -> str:
        """
        生成文本格式的报告
        
        Args:
            result: 分析结果
            
        Returns:
            str: 文本报告
        """
        lines = []
        lines.append("视频分析报告")
        lines.append("=" * 40)
        lines.append("")
        
        # 基本信息
        lines.append("【基本信息】")
        lines.append(f"视频ID: {result.get('video_id', '-')}")
        lines.append(f"分析时间: {result.get('analysis_time', '-')}")
        lines.append(f"状态: {result.get('status', 'completed')}")
        
        if "video_info" in result:
            video_info = result["video_info"]
            lines.append(f"视频时长: {video_info.get('duration', 0):.2f}秒")
            lines.append(f"分辨率: {video_info.get('width', 0)}x{video_info.get('height', 0)}")
            lines.append(f"帧率: {video_info.get('fps', 0):.2f}fps")
        
        lines.append(f"总体置信度: {result.get('overall_confidence', 0):.1f}%")
        lines.append("")
        
        # 编码类目分析结果
        lines.append("【编码类目分析结果】")
        lines.append("-" * 40)
        
        if "encoding_categories" in result:
            for category_name, category_data in result["encoding_categories"].items():
                lines.append(f"{category_name}:")
                lines.append(f"  - 选项: {category_data.get('option', '-')}")
                lines.append(f"  - 置信度: {category_data.get('confidence', 0)}%")
                if "reasoning" in category_data:
                    lines.append(f"  - 推理: {category_data['reasoning']}")
                lines.append("")
        else:
            lines.append("无编码类目分析结果")
            lines.append("")
        
        # 补充信息
        if "complementary_info" in result:
            lines.append("【补充信息】")
            lines.append("-" * 40)
            for key, value in result["complementary_info"].items():
                lines.append(f"{key}: {value}")
            lines.append("")
        
        # 备注
        lines.append("【备注】")
        lines.append("此报告由AI大模型视频编码分析系统自动生成")
        lines.append(datetime.now().strftime("生成时间: %Y-%m-%d %H:%M:%S"))
        
        return "\n".join(lines)
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        获取统计信息
        
        Returns:
            Dict[str, Any]: 统计数据
        """
        # 从结果管理器获取统计
        all_results = self.result_manager.list_all_results()
        
        # 计算统计数据
        status_count = {}
        confidences = []
        
        for result in all_results:
            # 统计状态
            status = result.get("status", "unknown")
            status_count[status] = status_count.get(status, 0) + 1
            
            # 收集置信度
            if "overall_confidence" in result and result["overall_confidence"] is not None:
                confidences.append(result["overall_confidence"])
        
        # 计算平均置信度
        avg_confidence = sum(confidences) / len(confidences) if confidences else 0
        
        return {
            "total_videos": len(all_results),
            "status_distribution": status_count,
            "average_confidence": avg_confidence,
            "analysis_count": len(all_results),
            "success_rate": status_count.get("completed", 0) / len(all_results) * 100 if all_results else 0
        }


# 创建全局实例
def get_vision_analysis_module() -> VisionAnalysisModule:
    """
    获取视觉分析模块实例
    
    Returns:
        VisionAnalysisModule: 视觉分析模块实例
    """
    global _vision_module
    
    try:
        if not hasattr(get_vision_analysis_module, '_vision_module'):
            # 从环境变量或配置文件加载API密钥
            api_key = os.environ.get('ZHIPU_API_KEY')
            
            # 尝试从.env文件加载
            if not api_key:
                env_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), '.env')
                if os.path.exists(env_path):
                    try:
                        with open(env_path, 'r', encoding='utf-8') as f:
                            for line in f:
                                if line.startswith('ZHIPU_API_KEY='):
                                    api_key = line.split('=', 1)[1].strip().strip('"')
                                    break
                    except:
                        pass
            
            # 创建模块实例
            get_vision_analysis_module._vision_module = VisionAnalysisModule(api_key=api_key)
        
        return get_vision_analysis_module._vision_module
    
    except Exception as e:
        logger.error(f"创建视觉分析模块实例失败: {str(e)}")
        raise


# 导出公共接口
def analyze_video(video_path: str, **kwargs) -> Dict[str, Any]:
    """
    分析视频的便捷函数
    
    Args:
        video_path: 视频文件路径
        **kwargs: 其他参数
        
    Returns:
        Dict[str, Any]: 分析结果
    """
    module = get_vision_analysis_module()
    return module.analyze_video(video_path, **kwargs)


def batch_analyze_videos(video_paths: List[str], **kwargs) -> Dict[str, Dict[str, Any]]:
    """
    批量分析视频的便捷函数
    
    Args:
        video_paths: 视频文件路径列表
        **kwargs: 其他参数
        
    Returns:
        Dict[str, Dict[str, Any]]: 分析结果字典
    """
    module = get_vision_analysis_module()
    return module.batch_analyze_videos(video_paths, **kwargs)


def generate_report(video_id: str, **kwargs) -> str:
    """
    生成报告的便捷函数
    
    Args:
        video_id: 视频ID
        **kwargs: 其他参数
        
    Returns:
        str: 报告内容
    """
    module = get_vision_analysis_module()
    return module.generate_analysis_report(video_id, **kwargs)


# 模块初始化代码
if __name__ == '__main__':
    try:
        print("视觉分析模块初始化中...")
        
        # 创建测试实例
        module = VisionAnalysisModule()
        
        print("\n视觉分析模块初始化成功!")
        print("\n可用功能:")
        print("- analyze_video(): 分析单个视频")
        print("- batch_analyze_videos(): 批量分析视频")
        print("- generate_analysis_report(): 生成分析报告")
        print("- export_all_results(): 导出所有结果")
        print("- get_statistics(): 获取统计信息")
        
        # 测试从环境变量加载API密钥
        has_api_key = module.analyzer._api_key is not None
        print(f"\nAPI密钥加载状态: {'已加载' if has_api_key else '未加载'}")
        
    except Exception as e:
        print(f"初始化失败: {str(e)}")
        sys.exit(1)